{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNh+WHHXHiw1v/WeZj1rcht",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Baijiong-Lin/LoRA-Torch/blob/main/examples/embedding.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "330HoZ7N_S3x",
        "outputId": "1d7f8b83-b37f-4b55-ad77-aa099b236083"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting git+https://github.com/microsoft/LoRA\n",
            "  Cloning https://github.com/microsoft/LoRA to /tmp/pip-req-build-qrvzzdp0\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/microsoft/LoRA /tmp/pip-req-build-qrvzzdp0\n",
            "  Resolved https://github.com/microsoft/LoRA to commit 998cfe4d351f4d6b4a47f0921dec2397aa0b9dfe\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Collecting git+https://github.com/Baijiong-Lin/LoRA-Torch\n",
            "  Cloning https://github.com/Baijiong-Lin/LoRA-Torch to /tmp/pip-req-build-ro2rj2g8\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/Baijiong-Lin/LoRA-Torch /tmp/pip-req-build-ro2rj2g8\n",
            "  Resolved https://github.com/Baijiong-Lin/LoRA-Torch to commit 34286e640a4d15dfe23c361f6d95b7cc55a8ec6b\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ],
      "source": [
        "!pip install git+https://github.com/microsoft/LoRA\n",
        "!pip install git+https://github.com/Baijiong-Lin/LoRA-Torch"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch, loralib, loratorch, copy\n",
        "import torch.nn as nn"
      ],
      "metadata": {
        "id": "cqadI5Fu_WUX"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model_lib = loralib.Embedding(10, 3, r=4)\n",
        "\n",
        "for k, v in model_lib.named_parameters():\n",
        "  print(k, v.size())\n",
        "\n",
        "\n",
        "loralib.mark_only_lora_as_trainable(model_lib)\n",
        "\n",
        "optimizer_lib = torch.optim.SGD(model_lib.parameters(), lr=0.1)\n",
        "\n",
        "for _ in range(3):\n",
        "    model_lib.train()\n",
        "\n",
        "    x = torch.rand(5, 10).long()\n",
        "\n",
        "    loss_lib = model_lib(x).sum()\n",
        "    optimizer_lib.zero_grad()\n",
        "    loss_lib.backward()\n",
        "    optimizer_lib.step()\n",
        "\n",
        "    x_test = torch.rand(*x.size()).long()\n",
        "    model_lib.eval()\n",
        "    print(model_lib(x_test).size())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "91GlkwcQD7YG",
        "outputId": "e87d2630-4ba5-456e-c8b3-63c25b36a9cb"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "weight torch.Size([10, 3])\n",
            "lora_A torch.Size([4, 10])\n",
            "lora_B torch.Size([3, 4])\n",
            "torch.Size([5, 10, 3])\n",
            "torch.Size([5, 10, 3])\n",
            "torch.Size([5, 10, 3])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model_torch = loratorch.Embedding(10, 3, r=4)\n",
        "\n",
        "for k, v in model_torch.named_parameters():\n",
        "  print(k, v.size())\n",
        "\n",
        "\n",
        "loratorch.mark_only_lora_as_trainable(model_torch)\n",
        "\n",
        "optimizer_torch = torch.optim.SGD(model_torch.parameters(), lr=0.1)\n",
        "\n",
        "for _ in range(3):\n",
        "    model_torch.train()\n",
        "\n",
        "    x = torch.rand(5, 10).long()\n",
        "\n",
        "    loss_torch = model_torch(x).sum()\n",
        "    optimizer_torch.zero_grad()\n",
        "    loss_torch.backward()\n",
        "    optimizer_torch.step()\n",
        "\n",
        "    x_test = torch.rand(*x.size()).long()\n",
        "    model_torch.eval()\n",
        "    print(model_torch(x_test).size())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NsLzmMXSC3RQ",
        "outputId": "48c1e55c-1236-42c5-ac87-c3d55e6bcf87"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "weight torch.Size([10, 3])\n",
            "w_lora_A torch.Size([4, 3])\n",
            "w_lora_B torch.Size([10, 4])\n",
            "torch.Size([5, 10, 3])\n",
            "torch.Size([5, 10, 3])\n",
            "torch.Size([5, 10, 3])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "q7yt9fq1GrnA"
      },
      "execution_count": 4,
      "outputs": []
    }
  ]
}